import json
import os
import subprocess
import sys
from openai import OpenAI

from src.modules.image_processor import ImageProcessor
from src.modules.ocr_processor import OCRProcessor
from src.modules.watermark_detector import WatermarkDetector
from src.utils.logger import log_error, log_info, log_warning
from src.utils.util import Util

class VideoProcessor:
    def __init__(self, api_key, status_callback=None):
        """初始化VideoProcessor并创建依赖的处理器"""
        if not api_key:
            raise ValueError("请提供有效的 OpenAI API 密钥")
        
        self.client = Util.get_openai_client(api_key=api_key)
        self.startupinfo = Util.get_startupinfo()
        self.status_callback = status_callback
        
        # 初始化各个处理器
        self.image_processor = ImageProcessor()
        self.ocr_processor = OCRProcessor()
        self.watermark_detector = WatermarkDetector()
    def update_status(self, message):
        """更新状态到UI"""
        if self.status_callback:
            self.status_callback(message)

    def classify_and_summarize_videos(self, video_data, max_retries=3):
        """
        调用 OpenAI API 对视频分类进行总结和简化，并合并相似分类
        增加重试机制处理token限制导致的视频漏分类问题
        :param video_data: 包含视频ID和分类的JSON数据
        :return: 合并后的分类结果（JSON格式）
        """
        log_info("开始分类和总结视频数据", ext=video_data)
        
        # 提取所有视频ID（排除"已有分类"键）
        video_ids = [vid_id for vid_id in video_data.keys() if vid_id != "已有分类"]
        self.update_status(f"正在处理 {len(video_ids)} 个视频文件...")

        retry_count = 0
        processed_result = None
        remaining_videos = dict(video_data)  # 初始时所有视频都需要处理

        while retry_count <= max_retries and remaining_videos:
            if retry_count > 0:
                self.update_status(f"第 {retry_count} 次重试分类，处理 {len(remaining_videos) - 1 if '已有分类' in remaining_videos else len(remaining_videos)} 个视频...")

            message_content = """
            给出的JSON数据中，key值是视频ID，视频ID是一个以mp4为后缀的文件名称字符串，value值是一个视频分类的数组。
            对分类进行总结和简化，用尽可能短但易于分辨的文字概括，一个视频ID只有一个分类。每个视频ID都要有分类。
            返回值中，视频ID必须完整，视频ID不会重复。
            有一个key值是 "已有分类", 对应的value是一个数组，数组中的元素是已经存在的分类。
            你需要将已有分类和新分类进行合并，返回的JSON数据中，key值是简化后的分类名称，value值是一个视频ID数组。
            将相似分类合并, 视频数据如下：
            """
            message_content += json.dumps(remaining_videos, ensure_ascii=False)
            message_content += "\n请严格按以下JSON格式返回结果：\n{\n  \"简化分类1\": [\"视频ID1\", \"视频ID2\"],\n  \"简化分类2\": [\"视频ID3\"],\n  ...\n}"

            messages = [
                {"role": "system", "content": "你是一个专业的分类和总结-助手。分类的结果必须是中文，且不能包含任何非中文字符。"},
                {"role": "user", "content": message_content}
            ]

            try:
                self.update_status(f"正在调用 OpenAI API 进行分类{' (重试)' if retry_count > 0 else ''}...")
                response = self.client.chat.completions.create(
                    model="gpt-4o-mini",
                    messages=messages,
                    max_tokens=16384,
                    temperature=0,
                    timeout=60
                )
                result = response.choices[0].message.content.replace("```json", "").replace("```", "").strip()
                # 验证分类是否为中文
                try:
                    classification_data = json.loads(result)
                    # Filter out empty arrays
                    classification_data = {k: v for k, v in classification_data.items() if v}
                    log_info("分类结果", ext=classification_data)
                    non_chinese_categories = []
                    
                    for category in classification_data.keys():
                        # 检查分类名称是否只包含中文字符
                        if not all('\u4e00' <= char <= '\u9fff' for char in category if not char.isspace()):
                            non_chinese_categories.append(category)
                    log_info("非中文分类", ext=non_chinese_categories)
                    # 如果存在非中文分类，请求修正
                    if non_chinese_categories:
                        self.update_status(f"检测到非中文分类，正在修正...")
                        
                        correction_messages = [
                            {"role": "system", "content": "你是一个专业的分类修正助手。你的任务是将所有非中文分类名称转换为纯中文。"},
                            {"role": "user", "content": f"""
                            请将以下JSON数据中的所有分类名称转换为纯中文，不要包含任何非中文字符。
                            需要修正的分类名称: {', '.join(non_chinese_categories)}
                            
                            原始JSON数据:
                            {result}
                            
                            请返回修正后的完整JSON数据，格式与原始数据相同。
                            """}
                        ]
                        
                        correction_response = self.client.chat.completions.create(
                            model="gpt-4o-mini",
                            messages=correction_messages,
                            max_tokens=16384,
                            temperature=0,
                            timeout=60
                        )
                        
                        result = correction_response.choices[0].message.content.replace("```json", "").replace("```", "").strip()
                        classification_data = json.loads(result)
                    
                    # 检查是否所有视频都已分类
                    classified_videos = set()
                    for videos in classification_data.values():
                        classified_videos.update(videos)
                    
                    # 找出未被分类的视频
                    current_video_ids = set(vid_id for vid_id in remaining_videos.keys() if vid_id != "已有分类")
                    missing_videos = current_video_ids - classified_videos
                    
                    if missing_videos:
                        self.update_status(f"发现 {len(missing_videos)} 个视频未被分类，准备重试...")
                        log_warning(f"未分类的视频: {', '.join(missing_videos)}")
                        
                        # 更新 processed_result
                        if processed_result is None:
                            processed_result = classification_data
                        else:
                            # 合并当前结果到已处理结果
                            for category, videos in classification_data.items():
                                if category in processed_result:
                                    processed_result[category].extend(
                                        [vid for vid in videos if vid not in processed_result[category]]
                                    )
                                else:
                                    processed_result[category] = videos
                        
                        # 为下一次重试准备只包含未分类视频的数据
                        new_remaining = {"已有分类": video_data.get("已有分类", [])}
                        for vid in missing_videos:
                            new_remaining[vid] = remaining_videos[vid]
                        
                        remaining_videos = new_remaining
                        retry_count += 1
                    else:
                        # 所有视频都已分类，合并结果并完成
                        if processed_result is None:
                            processed_result = classification_data
                        else:
                            for category, videos in classification_data.items():
                                if category in processed_result:
                                    processed_result[category].extend(
                                        [vid for vid in videos if vid not in processed_result[category]]
                                    )
                                else:
                                    processed_result[category] = videos
                        
                        remaining_videos = {}  # 清空剩余视频以结束循环
                    
                except json.JSONDecodeError as e:
                    error_msg = f"解析模型输出为JSON时出错: {e}"
                    self.update_status(error_msg)
                    log_error(error_msg)
                    retry_count += 1
                    
            except Exception as e:
                error_msg = f"调用 OpenAI API 时出错: {e}"
                self.update_status(error_msg)
                log_error(error_msg)
                retry_count += 1
        
        if not processed_result and retry_count > max_retries:
            self.update_status("分类失败，达到最大重试次数")
            return None
        
        # 检查是否有未分类的视频
        if remaining_videos and '已有分类' in remaining_videos:
            del remaining_videos['已有分类']
        
        if remaining_videos:
            self.update_status(f"警告：仍有 {len(remaining_videos)} 个视频未能分类")
        else:
            self.update_status("分类完成，所有视频已成功分类")
        
        return json.dumps(processed_result, ensure_ascii=False)

    def extract_frames(self, video_path, frame_count=20):
        """
        从视频中平均抽取指定数量的帧
        :param video_path: 视频文件路径
        :param frame_count: 要抽取的帧数（默认为5）
        :return: 抽取的帧图片路径列表
        """
        if not os.path.exists(video_path):
            raise FileNotFoundError(f"视频文件未找到: {video_path}")

        # 创建存储帧的临时目录
        frames_dir = video_path.replace(".mp4", "_frames")
        os.makedirs(frames_dir, exist_ok=True)

        try:
            # 获取视频的总时长
            # Construct the ffmpeg command to get video information in JSON format
            cmd = [
                "ffprobe", 
                "-v", "error",               # Only show errors
                "-print_format", "json",     # Output in JSON format
                "-show_format",              # Show format information
                "-show_streams",             # Show stream information
                video_path
            ]
            # Run the command
            result = subprocess.run(cmd,
                                    capture_output=True,
                                    text=True,
                                    encoding='utf-8',  # 防止 UnicodeDecodeError
                                    startupinfo=self.startupinfo
                                    )

            # Parse the JSON output to get duration
            probe = json.loads(result.stdout)
            duration = float(probe['format']['duration'])

            # 计算每帧的时间间隔
            interval = duration / frame_count

            log_info(f"视频时长: {duration:.2f}秒, 每帧间隔: {interval:.2f}秒")
            frame_paths = []
            for i in range(frame_count):
                timestamp = i * interval
                frame_path = os.path.join(frames_dir, f"frame_{i + 1}.jpg")
                
                # 构建ffmpeg命令
                # Construct command directly without ffmpeg Python library
                cmd = [
                    "ffmpeg", 
                    "-ss", str(timestamp), 
                    "-i", video_path, 
                    "-vframes", "1", 
                    "-f", "image2", 
                    "-c:v", "mjpeg", 
                    "-loglevel", "error",
                    "-y",  # Overwrite output file without asking
                    frame_path
                ]
                
                # Execute the command
                subprocess.run(cmd, startupinfo=self.startupinfo, stderr=subprocess.PIPE, stdout=subprocess.PIPE, encoding='utf-8')
                frame_paths.append(frame_path)

            log_info(f"成功抽取帧图片: {len(frame_paths)}个")
            self.image_processor.process_edge_images_in_folder(frames_dir)
            log_info(f"成功处理边缘图片: {len(frame_paths)}个")
            return frame_paths
        except Exception as e:
            raise RuntimeError(f"抽取帧图片时出错: {e}")






